package classifier;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Vector;

public class NaiveBayesClassifier extends GeneralClassifier<Vector<String>> {
	
	public NaiveBayesClassifier(HashSet<String> vocabulary)
	{
		LABEL_COUNT = 2;
		classPrior = new double[LABEL_COUNT];
		this.vocabulary = vocabulary;
	}

	public int classify(Vector<String> data, double level) 
	{
		double prC0 = 0;
		double prC1 = 0;
		double alpha = 1;
		
		// See lecture notes (slide 30)
		for(String s : data)
		{
			if(termCountPerClass.get(0).containsKey(s))
			{
				prC0 += Math.log(((double)termCountPerClass.get(0).get(s)+alpha)/((double)totalTermCountPerClass[0]+alpha*vocabulary.size()));
			}
			else
			{
				prC0 += Math.log(((double)0+alpha)/((double)totalTermCountPerClass[0]+alpha*vocabulary.size()));
			}
			if(termCountPerClass.get(1).containsKey(s))
			{
				prC1 += Math.log(((double)termCountPerClass.get(1).get(s)+alpha)/((double)totalTermCountPerClass[1]+alpha*vocabulary.size()));
			}
			else
			{
				prC1 += Math.log(((double)0+alpha)/((double)totalTermCountPerClass[1]+alpha*vocabulary.size()));
			}
		}
		prC0 += Math.log(classPrior[0]);
		prC1 += Math.log(classPrior[1]);
		
		
		double value = (prC0-prC1)/(Math.abs(prC0+prC1)); 
		if( value> level)
		{
			return 0;
		}
		else
		{
			return 1;
		}
	}

	/** 
	 * naive bayes training according to the lecture notes (slide 30)
	 */
	@Override
	public void train(Vector<LabeledData<Vector<String>>> trainingData) 
	{
		totalNumberOfDocuments = trainingData.size();
		initCounters();
		
		for(int label=0; label< LABEL_COUNT; label++)
		{
			int countSameClass = 0;
			for(LabeledData<Vector<String>> d : trainingData)
			{
				if(d.getLabel() == label)
				{
					addTermCountsForClass(d.getData(),d.getLabel());
					countSameClass++;
				}
			}
			// Assign the prior probability for class c_label Pr(c_label)
			classPrior[label] = ((double) countSameClass)/((double)totalNumberOfDocuments);
		}
		
	}
	
	private void initCounters()
	{
		// Throw away any old list of term counts or term per class counts
		totalTermCountPerClass = new int[LABEL_COUNT];
		termCountPerClass = new Vector<HashMap<String, Integer>>();
		for(int i=0; i< LABEL_COUNT; i++)
		{
			termCountPerClass.add(new HashMap<String, Integer>());
		}
		
	}
	
	private void addTermCountsForClass(Vector<String> data, int label)
	{
		for(String s : data)
		{
			if(vocabulary.contains(s))
			{
				if(termCountPerClass.get(label).containsKey(s))
				{
					termCountPerClass.get(label).put(s, termCountPerClass.get(label).get(s)+1);
				}
				else
				{
					termCountPerClass.get(label).put(s, 1);
				}
			}
			totalTermCountPerClass[label]++;
		}
		
	}

	private final int LABEL_COUNT;
	private int totalNumberOfDocuments;
	private double[] classPrior;
	private HashSet<String> vocabulary;
	
	/**
	 * Stores the number of occurrences for each term in the vocabulary
	 */
	private int [] totalTermCountPerClass;
	
	/**
	 * Stores the number of occurrences for each term in the vocabulary for each class separately
	 */
	private Vector<HashMap<String, Integer>> termCountPerClass;
	
}
